import os
from dotenv import load_dotenv
from typing import TypedDict

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph

# 加载 .env 文件中的环境变量
load_dotenv()

# 1. 定义状态 (State)
# 状态是一个 TypedDict，它定义了我们图在执行过程中会跟踪的数据结构。
# 这里我们只跟踪一个 "topic" 字符串。
class GraphState(TypedDict):
    """
    Represents the state of our graph.

    Attributes:
        topic: The topic to generate a joke about.
    """
    topic: str
    joke: str


# 2. 定义节点 (Nodes)
# 节点是执行工作的函数。它们接收当前的状态，并返回一个字典来更新状态。
def generate_joke(state: GraphState) -> dict[str, str]:
    """
    Generates a joke about the given topic.

    Args:
        state: The current graph state.

    Returns:
        A dictionary with the generated joke.
    """
    print("---GENERATING JOKE---")
    topic = state["topic"]
    
    # 我们使用 ChatOpenAI 来生成笑话
    # langchain-openai 会自动从环境变量加载 OPENAI_API_KEY 和 OPENAI_BASE_URL
    # 为了清晰，我们在这里显式传递 base_url
    llm = ChatOpenAI(
        model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL")
    )
    
    prompt = f"Tell me a short joke about {topic}"
    
    response = llm.invoke(prompt)
    
    joke = response.content
    print(f"---JOKE: {joke}---")

    return {"joke": joke}

# 3. 定义图 (Graph)
# StateGraph 是我们构建图的基础。
workflow = StateGraph(GraphState)

# 4. 添加节点和边 (Nodes and Edges)
# 我们将刚才定义的函数添加为图中的一个节点。
workflow.add_node("generate_joke", generate_joke)

# 设置入口点
workflow.set_entry_point("generate_joke")

# 设置出口点
workflow.set_finish_point("generate_joke")

# 5. 编译图 (Compile)
# 将图编译成一个可执行的 `Runnable` 对象。
app = workflow.compile()

# 6. 运行图 (Run)
if __name__ == "__main__":
    # 调用图，传入初始状态
    inputs = {"topic": "AI"}
    result = app.invoke(inputs)

    # 打印最终结果
    print("\n\nFinal Result:")
    print(result)
    print("\nGenerated Joke:")
    print(result["joke"]) 